import warnings
warnings.filterwarnings("ignore")

import os
import pandas as pd
import numpy as np
from keras.models import load_model
from tensorflow.keras import Input
from tensorflow.keras.models import Model
from tensorflow.keras.preprocessing.sequence import pad_sequences
from tensorflow.keras.layers import Dense, Embedding, LSTM, concatenate, Bidirectional,Dropout

model_path = '../../user_data/model_data/weights.best.hdf5'
prediction_path = '../../prediction_result/result.tsv'
df_test = pd.read_table('../../tcdata/oppo_breeno_round1_data/gaiic_track3_round1_testA_20210228.tsv', names=['q1', 'q2']).fillna("0")
# parameters
Vocab_size = 30000
Embedding_out_size = 32
LSTM_units = 32
Dense_units = 32


def toList(x):
    return list(map(int, x.split(" ")))


def preprocessing():
    max_length = 100
    df_test['q1'] = list(map(toList, df_test['q1'].tolist()[:]))
    df_test['q2'] = list(map(toList, df_test['q2'].tolist()[:]))

    test_sequence1 = pad_sequences(df_test['q1'], maxlen=max_length, padding='post')
    test_sequence2 = pad_sequences(df_test['q2'], maxlen=max_length, padding='post')
    return test_sequence1, test_sequence2


def test(test_sequence1, test_sequence2):
    text_input1 = Input(shape=(None,), dtype='int32')
    embedding1 = Embedding(Vocab_size, Embedding_out_size)(text_input1)
    encoded_text1 = Bidirectional(LSTM(LSTM_units))(embedding1)

    text_input2 = Input(shape=(None,), dtype='int32')
    embedding2 = Embedding(Vocab_size, Embedding_out_size)(text_input2)
    encoded_text2 = Bidirectional(LSTM(LSTM_units))(embedding2)
    concatenated = concatenate([encoded_text1, encoded_text2], axis=-1)

    output = Dense(Dense_units, activation='relu')(concatenated)
    output = Dense(1, activation='sigmoid')(output)

    model = Model([text_input1, text_input2], output)
    model.load_weights(model_path, by_name=True)

    prediction_out = model.predict([test_sequence1, test_sequence2])
    return prediction_out


def toResult(x):
    return float(x[0])


def run():
    test_sequence1, test_sequence2 = preprocessing()
    prediction_out = test(test_sequence1, test_sequence2)

    prediction_out = list(map(toResult, prediction_out))
    df = pd.DataFrame(prediction_out)
    df.to_csv(prediction_path, index=False, header=False, sep='\t')


if __name__ == '__main__':
    run()